import os
import requests
import zipfile
import threading
from tqdm import tqdm
from src.config import MODEL_PATHS
from src.config_manager import load_settings

MODEL_URLS = {
    'zh': 'https://alphacephei.com/vosk/models/vosk-model-cn-0.22.zip',
    'en': 'https://alphacephei.com/vosk/models/vosk-model-en-us-0.22.zip',
    'ja': 'https://alphacephei.com/vosk/models/vosk-model-ja-0.22.zip',
    'de': 'https://alphacephei.com/vosk/models/vosk-model-de-0.21.zip',
    'fr': 'https://alphacephei.com/vosk/models/vosk-model-fr-0.22.zip',
    'ru': 'https://alphacephei.com/vosk/models/vosk-model-ru-0.22.zip',
    'es': 'https://alphacephei.com/vosk/models/vosk-model-es-0.22.zip',
    'pt': 'https://alphacephei.com/vosk/models/vosk-model-pt-0.22.zip'
}

def download_model(lang_code):
    """
    下载并解压指定语言模型
    :param lang_code: 语言代码 (en, ja等)
    """
    settings = load_settings()
    proxies = {}
    if settings['proxy']['enable']:
        if settings['proxy']['http']:
            proxies['http'] = settings['proxy']['http']
        if settings['proxy']['https']:
            proxies['https'] = settings['proxy']['https']
        auth = None
        if settings['proxy']['username']:
            auth = (settings['proxy']['username'], settings['proxy']['password'])
            
    if lang_code not in MODEL_URLS:
        raise ValueError(f"不支持的语言代码: {lang_code}")
    url = MODEL_URLS[lang_code]
    
    model_zip = f"temp_{lang_code}.zip"
    model_path = MODEL_PATHS[lang_code]
    
    # 创建模型目录
    os.makedirs(os.path.dirname(model_path), exist_ok=True)
    
    try:
        # 下载文件
        print(f"开始下载 {lang_code} 模型...")
        response = requests.get(url, stream=True)
        total_size = int(response.headers.get('content-length', 0))
        
        with open(model_zip, 'wb') as f, tqdm(
            desc=f"下载 {lang_code}",
            total=total_size,
            unit='iB',
            unit_scale=True,
            unit_divisor=1024,
        ) as bar:
            for data in response.iter_content(chunk_size=1024):
                if not getattr(threading.current_thread(), "do_run", True):
                    raise Exception("用户取消下载")
                size = f.write(data)
                bar.update(size)
        
        # 解压文件
        print(f"解压 {lang_code} 模型...")
        with zipfile.ZipFile(model_zip, 'r') as zip_ref:
            zip_ref.extractall(os.path.dirname(model_path))
        
        # 清理临时文件
        os.remove(model_zip)
        return True
    
    except Exception as e:
        print(f"下载失败: {str(e)}")
        if os.path.exists(model_zip):
            os.remove(model_zip)
        return False 

def check_model_installed(lang_code):
    """检查模型是否已安装"""
    return os.path.exists(MODEL_PATHS.get(lang_code, "")) 